local redisKey = KEYS[1] -- 要进行限流的Key
local consume_tokens = tonumber(ARGV[1])  -- 每次请求消耗的令牌数
local curr_mills = tonumber(ARGV[2]) -- 当前时间毫秒值
local max_bucket_cap = tonumber(ARGV[3]) -- 最大桶数
local period = tonumber(ARGV[4]) -- 限流时间周期(秒)
local generate_rate = max_bucket_cap / period  --每秒产生的令牌数
local expire_seconds = 60 * 10  -- key存活时间,10分钟
local result = 1  -- 返回结果 1:成功, 0:失败

-- 查询上次请求的时间和剩余令牌数量
local limiter_info = redis.pcall("HMGET", redisKey, "last_time", "remaining_tokens")

-- 无数据,则表明当前为第一次请求
if not limiter_info[1] or not limiter_info[2] then
    local remaining_tokens = max_bucket_cap - consume_tokens  -- 剩余令牌数
    local last_time = curr_mills + period * 1000 --令牌最后期限
    redis.pcall("HMSET", redisKey,
            "remaining_tokens", remaining_tokens,
            "last_time", last_time,
            "max_bucket_cap", max_bucket_cap,
            "generate_rate", generate_rate,
            "period", period)
    redis.call("expire", redisKey, expire_seconds)
    return result
end

-- 执行至此, 表明不是第一次请求, 已创建令牌桶
local last_time = tonumber(limiter_info[1])   -- 令牌最后有效期限
local remaining_tokens = tonumber(limiter_info[2])  -- 上次请求最终剩余桶数
local current_remaining_tokens = 0  -- 计算当前剩余桶数

-- 非初次请求, 核对令牌数
if last_time > 0 then
    -- 新产生令牌数 = 距离上次请求的秒数 * 令牌产生速率
    local new_permits = math.floor((curr_mills - last_time) / 1000 * generate_rate)

    if new_permits <= 0 then
        new_permits = 0 -- 新增令牌小于0, 即还在初始令牌最后期限之内
    end

    current_remaining_tokens = new_permits + remaining_tokens   -- 当前剩余 = 新产生 + 上次剩余
    current_remaining_tokens = math.min(current_remaining_tokens, max_bucket_cap) -- 当前所有桶数 不能超过 最大桶数
end

if current_remaining_tokens >= consume_tokens then
    current_remaining_tokens = current_remaining_tokens - consume_tokens
else
    result = 0  -- 当前令牌桶数不够消耗,请求失败
end

if curr_mills > last_time then
    last_time = curr_mills -- 刷新令牌桶有效期
end

redis.pcall("HMSET", redisKey,
        "remaining_tokens", current_remaining_tokens,
        "last_time", last_time,
        "max_bucket_cap", max_bucket_cap,
        "generate_rate", generate_rate,
        "period", period)

redis.call("expire", redisKey, expire_seconds)
return result -- 返回结果 1:成功, 0:失败